# -*- coding: utf-8 -*-

import cv2
import numpy as np
from djitellopy import Tello
tello = Tello()
tello.connect()
tello.streamon()
print(tello.get_battery())

def detect_circles(camera_index=0):
    cap = cv2.VideoCapture(camera_index)
    forward_count = 0

    while True:
        ret, frame = cap.read()
        frame = tello.get_frame_read().frame
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        # ⭐ 对比度增强
        gray = cv2.equalizeHist(gray)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        gray = clahe.apply(gray)

        # ⭐ 去噪 + 锐化
        blurred = cv2.GaussianBlur(gray, (3, 3), 0)
        sharp = cv2.Laplacian(blurred, cv2.CV_64F)
        sharp = np.uint8(np.absolute(sharp))

        # ⭐ Canny 边缘检测

        edges = cv2.Canny(sharp, 50, 125)


        circles = cv2.HoughCircles(edges, cv2.HOUGH_GRADIENT, dp=1, minDist=100,
                                   param1=50, param2=150, minRadius=10, maxRadius=200)

        if circles is not None:
            circles = np.round(circles[0, :]).astype("int")
            for (x, y, r) in circles:
                cv2.circle(frame, (x, y), r, (0, 255, 0), 4)

                mask = np.zeros_like(gray)
                cv2.circle(mask, (x, y), r, (255, 255, 255), -1)
                contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
                if contours:
                    M = cv2.moments(contours[0])
                    center_x = int(M["m10"] / M["m00"])
                    center_y = int(M["m01"] / M["m00"])

                    if center_x < frame.shape[1] // 2:
                        print("Fly left")
                    elif center_x > frame.shape[1] // 2:
                        print("Fly right")
                    if center_y < frame.shape[0] // 2:
                        print("Fly up")
                    elif center_y > frame.shape[0] // 2:
                        print("Fly down")

        forward_count += 1
        print(f"Move forward ({forward_count})")

        cv2.imshow('Original Frame', frame)
        cv2.imshow('Grayscale Image', gray)
        cv2.imshow('Canny Edge Detection', edges)
        cv2.imshow("Sharp", sharp)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


# apply
detect_circles(camera_index=0)